# 1. 导入核心模块
from qwen_agent.agents import Assistant
import os
from dotenv import load_dotenv  # 用于加载环境变量（需安装：pip install python-dotenv）

import logging

# 仅保留ERROR级别以上的日志（屏蔽INFO日志）
logging.basicConfig(level=logging.ERROR)
# 若需保留Qwen-Agent的关键日志，可单独设置其日志级别
logging.getLogger("qwen_agent").setLevel(logging.WARNING)

# 2. 加载环境变量（避免API密钥硬编码）
load_dotenv()
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    raise ValueError("请先设置DASHSCOPE_API_KEY环境变量")

# 3. 配置LLM（使用通义千问模型，适配表格解析需求）
llm_cfg = {
    "model": "qwen-plus",  # 推荐qwen-max/qwen-plus（支持复杂表格理解）
    "model_type": "qwen_dashscope",  # 固定值，指定调用DashScope服务
    "api_key": api_key,
    "stream": True,  # 启用流式输出（适合实时展示答案）
    "temperature": 0.1  # 低温度→答案更精准（考核规则为固定条款，无需创造性）
}

# 4. 配置RAG参数（针对考核办法的表格+条款特点）
rag_cfg = {
    "top_k": 4,  # 每次检索返回4条相关知识（避免漏检表格行/条款）
    "chunk_size": 500,  # 分块大小：适配PDF章节长度（如单条考核规则约200-300字）
    "table_parse": True  # 强制启用表格解析（关键！确保第四章表格数据被正确提取）
}

# 5. 初始化RAG Agent，加载考核办法PDF
# 注意：files参数为本地PDF路径（相对/绝对路径均可）
agent = Assistant(
    llm=llm_cfg,
    function_list=[],  # 纯RAG场景，无需额外工具
    name="浦发个金客户经理考核查询助手",
    description="专注于解答《上海浦东发展银行西安分行个金客户经理管理考核暂行办法》的相关问题，包括职位标准、考核分值、扣分规则等",
    files=["./浦发上海浦东发展银行西安分行个金客户经理考核办法.pdf"],  # 替换为你的PDF路径
    rag_cfg=rag_cfg
)

print("PDF文档加载成功！可开始查询考核办法相关问题。")


# 1. 定义查询函数（支持流式输出+表格格式化）
def query_assessment_rule(question: str) -> None:
    print(f"\n用户查询：{question}")
    print("=" * 50)
    print("助手回答：", end="")

    # 构造用户消息（确保格式正确）
    messages = [
        {
            "role": "user",
            "content": question  # 纯文本查询，若带文件需用列表格式
        }
    ]

    # 关键：启用流式输出并完整捕获回答
    full_answer = ""
    printed_content = ""  # 记录已经打印的内容，避免重复打印
    try:
        # 调用Agent时确保llm配置中stream=True（此前配置已包含）
        response_count = 0
        for response in agent.run(messages=messages, stream=True):
            response_count += 1
            # Qwen-Agent的流式响应格式：response -> "choices" -> [0] -> "delta" -> "content"
            content = ""
            if isinstance(response, dict):
                if "choices" in response and len(response["choices"]) > 0:
                    delta = response["choices"][0].get("delta", {})
                    content = delta.get("content", "")
                elif "message" in response:
                    # 处理非流式响应
                    content = response["message"].get("content", "")
                else:
                    # 尝试直接获取内容
                    content = response.get("content", "") or str(response)
            elif isinstance(response, str):
                # 处理可能的字符串响应
                content = response
            elif isinstance(response, list) and len(response) > 0:
                # 处理列表响应，获取其中的内容
                first_item = response[0]
                if isinstance(first_item, dict):
                    if "content" in first_item:
                        content = first_item["content"]
                    elif "text" in first_item:
                        content = first_item["text"]
            
            # 只打印新内容，避免重复
            if content and content not in printed_content:
                # 检查是否是完整内容的扩展
                if content.startswith(printed_content) and len(content) > len(printed_content):
                    new_content = content[len(printed_content):]
                    full_answer = content
                    printed_content = content
                    print(new_content, end="", flush=True)
                else:
                    full_answer += content
                    printed_content += content
                    print(content, end="", flush=True)
        
        # 如果没有获取到流式内容，尝试获取完整响应
        if not full_answer:
            print("\n[警告] 流式响应未获取到内容，尝试非流式响应...")
            responses = list(agent.run(messages=messages, stream=False))
            if responses:
                for response in responses:
                    content = ""
                    if isinstance(response, dict):
                        if "choices" in response and len(response["choices"]) > 0:
                            content = response["choices"][0].get("message", {}).get("content", "")
                        else:
                            content = response.get("content", "") or response.get("text", "") or str(response)
                    elif isinstance(response, str):
                        content = response
                    elif isinstance(response, list) and len(response) > 0:
                        first_item = response[0]
                        if isinstance(first_item, dict):
                            content = first_item.get("content", "") or first_item.get("text", "")
                    
                    if content and content not in full_answer:
                        full_answer += content
                        print(content, end="", flush=True)
                        
    except Exception as e:
        print(f"\n查询过程中出现错误: {e}")
        full_answer = "抱歉，查询过程中出现错误，请稍后重试。"

    # 打印完整答案（可选，用于调试）
    print(f"\n\n完整答案：{full_answer}")
    
    # 4. 若答案包含表格数据，用Pandas格式化展示（可选）
    try:
        import re
        import pandas as pd
        # 匹配PDF中的表格数据（如"客户经理助理 | 5 | 90 | 300万 | | 500张"）
        table_pattern = r"([^|]+?)\s*\|\s*([^|]+?)\s*\|\s*([^|]+?)\s*\|\s*([^|]+?)\s*\|\s*([^|]*?)\s*\|\s*([^|]+?)"
        table_matches = re.findall(table_pattern, full_answer)
        if table_matches:
            print("\n\n格式化表格：")
            columns = ["职位", "行员级别", "考核分值", "储蓄业务（万）", "个贷业务（万）", "卡业务（张）"]
            df = pd.DataFrame(table_matches, columns=columns)
            print(df.to_string(index=False))
    except ImportError:
        print("\n\n注意：需要安装pandas库才能显示格式化表格 (pip install pandas)")
    except Exception as e:
        print(f"\n[警告] 表格解析过程中出现错误: {e}")


# 2. 测试典型查询案例
if __name__ == "__main__":
    # 案例1：查询职位准入标准
    query_assessment_rule("客户经理助理5级的储蓄业务准入标准是多少？卡业务需要多少张有效卡？")

    # 案例2：查询考核分值
    #query_assessment_rule("高级客户经理1级的考核分值是多少？资深客户经理5级呢？")

    # 案例3：查询工作质量扣分规则
   # query_assessment_rule("个金客户经理出现跨月逾期，单笔不超过10万元且当季收回，扣多少分？廉洁自律问题扣多少分？")

    # 案例4：查询折算规则
    #query_assessment_rule("考核办法中，储蓄、个贷、有效卡的折算标准是什么？50万储蓄相当于多少分？")